#include <iostream>
#include <opencv2/opencv.hpp>

using namespace std;
using namespace cv;

//创建一个灯条的类
class LightDescriptor
{
public:
    float width, length, angle, area;
    Point2f center;

public:
    LightDescriptor(){};

    LightDescriptor(const cv::RotatedRect &light)
    {
        width = light.size.width;//宽
        length = light.size.height;//长
        center = light.center;//坐标点
        angle = light.angle;//旋转角度
        area = light.size.area();//面积
    }
};

double compareContours(const vector<Point>& contour1, const vector<Point>& contour2)//匹配轮廓，识别贴纸的函数
{
    double result = matchShapes(contour1, contour2, CONTOURS_MATCH_I1, 0);
    return result;
}

int main()
{
    VideoCapture video;
    video.open("/home/lwj/opencv.21/armor_blue.avi");//读取视频

    if (!video.isOpened())
    {
        cout << "Error opening video file" << endl;
        return -1;
    }
    //初始化灯条识别的参数
    Mat frame, hsv, binary, Gaussian, dilatee;
    Mat element = getStructuringElement(MORPH_RECT, Size(5, 5));
    vector<vector<Point>> contours;//存放像素坐标
    vector<Vec4i> hierarchy;//存放轮廓
    double prevTime = 0;
    double currentTime = 0;
    double timeInterval = 0;
    double prevX = 0;
    double prevY = 0;
    double currentX = 0;
    double currentY = 0;
    double deltaX = 0;
    double deltaY = 0;
    double velocityX = 0;
    double velocityY = 0;

    // 读取视频的第一帧作为参考图像
    Mat referenceImage;
    video >> referenceImage;

    // 将参考图像转为灰度图像
    Mat referenceGray;
    cvtColor(referenceImage, referenceGray, COLOR_BGR2GRAY);

    // 检测参考图像的轮廓
    vector<vector<Point>> referenceContours;
    vector<Vec4i> hierar;
    findContours(referenceGray, referenceContours, hierar, RETR_EXTERNAL, CHAIN_APPROX_SIMPLE);

    // 输入三张图片
    Mat inputImage1 = imread("/home/lwj/opencv.21/4.jpg");
    Mat inputImage2 = imread("/home/lwj/opencv.21/pao.jpg");
    Mat inputImage3 = imread("/home/lwj/opencv.21/tai.jpg");

    // 将输入图像转为灰度图像
    Mat inputGray1, inputGray2, inputGray3;
    cvtColor(inputImage1, inputGray1, COLOR_BGR2GRAY);
    cvtColor(inputImage2, inputGray2, COLOR_BGR2GRAY);
    cvtColor(inputImage3, inputGray3, COLOR_BGR2GRAY);

    // 检测输入图像的轮廓
    vector<vector<Point>> inputContours1, inputContours2, inputContours3;
    findContours(inputGray1, inputContours1, hierar, RETR_EXTERNAL, CHAIN_APPROX_SIMPLE);
    findContours(inputGray2, inputContours2, hierar, RETR_EXTERNAL, CHAIN_APPROX_SIMPLE);
    findContours(inputGray3, inputContours3, hierar, RETR_EXTERNAL, CHAIN_APPROX_SIMPLE);

    // 计算每个输入图像与参考图像轮廓的匹配度
    double match1 = compareContours(referenceContours[0], inputContours1[0]);
    double match2 = compareContours(referenceContours[0], inputContours2[0]);
    double match3 = compareContours(referenceContours[0], inputContours3[0]);

    for (;;)
    {
        double t = (double)getTickCount();
        video >> frame;

        if (frame.empty())
        {
            break;
        }

        cvtColor(frame, hsv, COLOR_BGR2HSV); // 转换为HSV颜色空间

        // 颜色过滤，只保留蓝色灯条的区域
        inRange(hsv, Scalar(113, 100, 100), Scalar(130, 255, 255), binary);

        GaussianBlur(binary, Gaussian, Size(5, 5), 0);//高斯滤波
        dilate(Gaussian, dilatee, element);//膨胀
        findContours(dilatee, contours, hierarchy, RETR_TREE, CHAIN_APPROX_NONE);//检测轮廓

        vector<LightDescriptor> lightInfos;

        for (int i = 0; i < contours.size(); i++)
        {
            double area = contourArea(contours[i]);
            if (area < 800 || area > 2000) // 添加轮廓面积限制条件
                {
                    continue;
                }
            if (area < 5 || contours[i].size() <= 1)
                {
                continue;
                }
            RotatedRect Light_Rec = fitEllipse(contours[i]);
            if (Light_Rec.size.width / Light_Rec.size.height > 4)
               { 
                continue;
               }
            lightInfos.push_back(LightDescriptor(Light_Rec));
        }

        for (size_t i = 0; i < lightInfos.size(); i++)//筛选灯条的参数
        {
            for (size_t j = i + 1; (j < lightInfos.size()); j++)
            {
                LightDescriptor &leftLight = lightInfos[i];
                LightDescriptor &rightLight = lightInfos[j];
                float angleGap_ = abs(leftLight.angle - rightLight.angle);
                float LenGap_ratio = abs(leftLight.length - rightLight.length) / max(leftLight.length, rightLight.length);
                float dis = pow(pow((leftLight.center.x - rightLight.center.x), 2) + pow((leftLight.center.y - rightLight.center.y), 2), 0.5);
                float meanLen = (leftLight.length + rightLight.length) / 2;
                float lengap_ratio = abs(leftLight.length - rightLight.length) / meanLen;
                float yGap = abs(leftLight.center.y - rightLight.center.y);
                float yGap_ratio = yGap / meanLen;
                float xGap = abs(leftLight.center.x - rightLight.center.x);
                float xGap_ratio = xGap / meanLen;
                float ratio = dis / meanLen;
                //筛选灯条的旋转矩形
                if (angleGap_ > 15 ||//矩形的角度
                    LenGap_ratio > 1.0 ||//矩形的长度差
                    lengap_ratio > 0.8 ||
                    yGap_ratio > 1.5 ||//旋转矩形与y轴的差
                    xGap_ratio > 2.2 ||//旋转矩形与x轴的差
                    xGap_ratio < 0.8 ||
                    ratio > 3 ||//面积和长度比例
                    ratio < 0.8)
                {
                    continue;
                }

                Point center = Point((leftLight.center.x + rightLight.center.x) / 2, (leftLight.center.y + rightLight.center.y) / 2);//计算灯条坐标
                RotatedRect rect = RotatedRect(center, Size(dis, meanLen), (leftLight.angle + rightLight.angle) / 2);
                Point2f vertices[4];
                rect.points(vertices);

                for (int i = 0; i < 4; i++)//绘制矩形框住装甲板
                {
                    line(frame, vertices[i], vertices[(i + 1) % 4], Scalar(0, 255, 0), 2.2);
                    line(frame, vertices[i], vertices[(i + 2) % 4], Scalar(0, 255, 0), 2.2);
                }
                circle(frame, (vertices[0] + vertices[2]) / 2, 5, Scalar(0, 0, 255), -1);//选取中心点标注颜色

                currentTime = (double)cv::getTickCount() / cv::getTickFrequency();
                timeInterval = currentTime - prevTime;
                prevTime = currentTime;
                //测速思想 对比前一帧图像的位置，分别求取在x和y轴上的位移距离，再除以时间，从而求出瞬时的二维速度
                currentX = center.x;
                currentY = center.y;
                deltaX = currentX - prevX;
                deltaY = currentY - prevY;
                velocityX = deltaX / timeInterval;
                velocityY = deltaY / timeInterval;
                //输出二维转速
                cout << "Velocity X: " << velocityX << " pixels/second" << endl;
                cout << "Velocity Y: " << velocityY << " pixels/second" << endl;

                prevX = currentX;
                prevY = currentY;

                putText(frame, "Blue", center, FONT_HERSHEY_TRIPLEX, 1.0, Scalar(255, 0, 0), 2);
            }
        }

        namedWindow("video", WINDOW_FREERATIO);
        imshow("video", frame);
        waitKey(1);
    }
    // 输出匹配度最高的图像
    if (match1 < match2 && match1 < match3)
     {
        imshow("识别的贴纸", inputImage1);
    } else if (match2 < match1 && match2 < match3) 
    {
        imshow("识别的贴纸", inputImage2);
    } else 
    {
        imshow("识别的贴纸", inputImage3);
    }
    waitKey(0);
    video.release();//释放视频
    destroyAllWindows();
    return 0;
}